# 练习使用 Sobel算子

import numpy as np
import cv2
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import pickle

# 读取图象
image = mpimg.imread('signs_vehicles_xygrad.png')
    
def mag_thresh(img, sobel_kernel=3, mag_thresh=(0, 255)):
    
    # 输出灰度图
    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    
    # 使用 Sobel 算子来计算 x y 得梯度
    sobelx = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=sobel_kernel)
    sobely = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=sobel_kernel)
    
    # 计算梯度大小
    gradmag = np.sqrt(sobelx**2 + sobely**2)

    # 转为8位
    scale_factor = np.max(gradmag)/255 
    gradmag = (gradmag/scale_factor).astype(np.uint8) 

    # 符合阈值条件 作为二值输出
    binary_output = np.zeros_like(gradmag)
    binary_output[(gradmag >= mag_thresh[0]) & (gradmag <= mag_thresh[1])] = 1

    #返回二值图象
    return binary_output
    
    
# 执行这个函数
mag_binary = mag_thresh(image, sobel_kernel=3, mag_thresh=(30, 100))

# 可视化结果
f, (ax1, ax2) = plt.subplots(1, 2, figsize=(24, 9))
f.tight_layout()
ax1.imshow(image)
ax1.set_title('Original Image', fontsize=50)
ax2.imshow(mag_binary, cmap='gray')
ax2.set_title('Thresholded Gradient', fontsize=50)
plt.subplots_adjust(left=0., right=1, top=0.9, bottom=0.)
plt.show()

